import os
import pinecone
import logging
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from collections import Counter
import numpy as np
import tqdm
from pinecone import Pinecone, ServerlessSpec

# 使用 Pinecone API 密钥创建客户端实例
pc = Pinecone("0b996cab-2029-4ed3-8646-20e9dec53ae2")
# 索引名称
index_name = "mnist-index"

# 获取现有索引列表
existing_indexes = pc.list_indexes()

# 检查索引是否存在，如果存在就删除
# 这个if是否需要，要看情况而定
# 比如有的时候，如果不想要重复删除在创建，这个if就可以不要
if any(index['name'] == index_name for index in existing_indexes):
    print(f"索引 '{index_name}' 已存在，正在删除...")
    pc.delete_index(index_name)
    print(f"索引 '{index_name}' 已成功删除。")
else:
    print(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
print(f"正在创建新索引 '{index_name}'...")
pc.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
print(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pc.Index(index_name)
print(f"已成功连接到索引 '{index_name}'。")

# 日志配置
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S')

# 加载 MNIST 数据
digits = load_digits(n_class=10)
X = digits.data
y = digits.target

# 分割数据集 (80% 训练, 20% 测试)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 准备上传数据
vectors = []
for i in range(len(X_train)):
    vector_id = str(i)
    vector_values = X_train[i].tolist()
    metadata = {"label": int(y_train[i])}
    vectors.append((vector_id, vector_values, metadata))

# 上传数据到 Pinecone 索引，显示进度条
batch_size = 100
logging.info("开始上传数据到 Pinecone 索引...")
for i in tqdm.tqdm(range(0, len(vectors), batch_size), total=int(len(vectors)/batch_size)):
    batch = vectors[i:i + batch_size]
    index.upsert(batch)

logging.info(f"数据上传完成，总共上传了 {len(vectors)} 条数据")

# 记录上传的数量
logging.info(f"成功创建索引，并上传了 {len(vectors)} 条数据")

# 开始测试, 使用20%的测试数据
correct_predictions = 0
logging.info("开始测试...")

for i in tqdm.tqdm(range(len(X_test)), total=len(X_test)):
    query_vector = X_test[i].tolist()
    results = index.query(vector=query_vector, top_k=11, include_metadata=True)
    labels = [match['metadata']['label'] for match in results['matches']]
    
    # 投票机制确定预测结果
    predicted_label = Counter(labels).most_common(1)[0][0]
    
    # 判断预测是否正确
    if predicted_label == y_test[i]:
        correct_predictions += 1

# 计算准确率
accuracy = correct_predictions / len(X_test)
logging.info(f"当 k=11 时，使用 Pinecone 的准确率为: {accuracy:.4f}")

# 根据需求决定是否删除索引
# pc.delete_index(index_name)

